{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/workflow/JSONalyze_query_engine.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# JSONalyze Query Engine\n",
    "\n",
    "JSONalyze Query Engine is designed to be wired typically after a calling(by agent, etc) of APIs, where we have the returned value as bulk instances of rows, and the next step is to perform statistical analysis on the data.\n",
    "\n",
    "With JSONalyze, under the hood, in-memory SQLite table is created with the JSON List loaded, the query engine is able to perform SQL queries on the data, and return the Query Result as answer to the analytical questions.\n",
    "\n",
    "This notebook walks through implementation of JSON Analyze Query Engine, using Workflows.\n",
    "\n",
    "Specifically we will implement [JSONalyzeQueryEngine](https://github.com/run-llama/llama_index/blob/main/docs/examples/query_engine/JSONalyze_query_engine.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since workflows are async first, this all runs fine in a notebook. If you were running in your own code, you would want to use `asyncio.run()` to start an async event loop if one isn't already running.\n",
    "\n",
    "```python\n",
    "async def main():\n",
    "    <async code>\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    import asyncio\n",
    "    asyncio.run(main())\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Workflow\n",
    "\n",
    "`jsonalyzer:`\n",
    "\n",
    "1. It takes a StartEvent as input and returns a JsonAnalyzerEvent.\n",
    "2. The function sets up an in-memory SQLite database, loads JSON data, generates a SQL query based on query using a LLM, executes the query, and returns the results along with the SQL query and table schema.\n",
    "\n",
    "`synthesize:`\n",
    "\n",
    "The function uses a LLM to synthesize a response based on the SQL query, table schema, and query results.\n",
    "\n",
    "The steps will use the built-in `StartEvent` and `StopEvent` events."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Event"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import Event\n",
    "from typing import Dict, List, Any\n",
    "\n",
    "\n",
    "class JsonAnalyzerEvent(Event):\n",
    "    \"\"\"\n",
    "    Event containing results of JSON analysis.\n",
    "\n",
    "    Attributes:\n",
    "        sql_query (str): The generated SQL query.\n",
    "        table_schema (Dict[str, Any]): Schema of the analyzed table.\n",
    "        results (List[Dict[str, Any]]): Query execution results.\n",
    "    \"\"\"\n",
    "\n",
    "    sql_query: str\n",
    "    table_schema: Dict[str, Any]\n",
    "    results: List[Dict[str, Any]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prompt Templates\n",
    "\n",
    "Here we define default `DEFAULT_RESPONSE_SYNTHESIS_PROMPT_TMPL`, `DEFAULT_RESPONSE_SYNTHESIS_PROMPT` and `DEFAULT_TABLE_NAME`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.prompts.prompt_type import PromptType\n",
    "from llama_index.core.prompts import PromptTemplate\n",
    "\n",
    "DEFAULT_RESPONSE_SYNTHESIS_PROMPT_TMPL = (\n",
    "    \"Given a query, synthesize a response based on SQL query results\"\n",
    "    \" to satisfy the query. Only include details that are relevant to\"\n",
    "    \" the query. If you don't know the answer, then say that.\\n\"\n",
    "    \"SQL Query: {sql_query}\\n\"\n",
    "    \"Table Schema: {table_schema}\\n\"\n",
    "    \"SQL Response: {sql_response}\\n\"\n",
    "    \"Query: {query_str}\\n\"\n",
    "    \"Response: \"\n",
    ")\n",
    "\n",
    "DEFAULT_RESPONSE_SYNTHESIS_PROMPT = PromptTemplate(\n",
    "    DEFAULT_RESPONSE_SYNTHESIS_PROMPT_TMPL,\n",
    "    prompt_type=PromptType.SQL_RESPONSE_SYNTHESIS,\n",
    ")\n",
    "\n",
    "DEFAULT_TABLE_NAME = \"items\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Workflow Itself\n",
    "\n",
    "With our events defined, we can construct our workflow and steps. \n",
    "\n",
    "Note that the workflow automatically validates itself using type annotations, so the type annotations on our steps are very helpful!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.base.response.schema import Response\n",
    "from llama_index.core.indices.struct_store.sql_retriever import (\n",
    "    DefaultSQLParser,\n",
    ")\n",
    "from llama_index.core.prompts.default_prompts import DEFAULT_JSONALYZE_PROMPT\n",
    "from llama_index.core.utils import print_text\n",
    "\n",
    "from llama_index.core.workflow import (\n",
    "    Context,\n",
    "    Workflow,\n",
    "    StartEvent,\n",
    "    StopEvent,\n",
    "    step,\n",
    ")\n",
    "\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "from IPython.display import Markdown, display\n",
    "\n",
    "\n",
    "class JSONAnalyzeQueryEngineWorkflow(Workflow):\n",
    "    @step\n",
    "    async def jsonalyzer(\n",
    "        self, ctx: Context, ev: StartEvent\n",
    "    ) -> JsonAnalyzerEvent:\n",
    "        \"\"\"\n",
    "        Analyze JSON data using a SQL-like query approach.\n",
    "\n",
    "        This asynchronous method sets up an in-memory SQLite database, loads JSON data,\n",
    "        generates a SQL query based on a natural language question, executes the query,\n",
    "        and returns the results.\n",
    "\n",
    "        Args:\n",
    "            ctx (Context): The context object for storing data during execution.\n",
    "            ev (StartEvent): The event object containing input parameters.\n",
    "\n",
    "        Returns:\n",
    "            JsonAnalyzerEvent: An event object containing the SQL query, table schema, and query results.\n",
    "\n",
    "        The method performs the following steps:\n",
    "        1. Imports the required 'sqlite-utils' package.\n",
    "        2. Extracts necessary data from the input event.\n",
    "        3. Sets up an in-memory SQLite database and loads the JSON data.\n",
    "        4. Generates a SQL query using a LLM based on the input question.\n",
    "        5. Executes the SQL query and retrieves the results.\n",
    "        6. Returns the results along with the SQL query and table schema.\n",
    "\n",
    "        Note:\n",
    "            This method requires the 'sqlite-utils' package to be installed.\n",
    "        \"\"\"\n",
    "        try:\n",
    "            import sqlite_utils\n",
    "        except ImportError as exc:\n",
    "            IMPORT_ERROR_MSG = (\n",
    "                \"sqlite-utils is needed to use this Query Engine:\\n\"\n",
    "                \"pip install sqlite-utils\"\n",
    "            )\n",
    "\n",
    "            raise ImportError(IMPORT_ERROR_MSG) from exc\n",
    "\n",
    "        await ctx.store.set(\"query\", ev.get(\"query\"))\n",
    "        await ctx.store.set(\"llm\", ev.get(\"llm\"))\n",
    "\n",
    "        query = ev.get(\"query\")\n",
    "        table_name = ev.get(\"table_name\")\n",
    "        list_of_dict = ev.get(\"list_of_dict\")\n",
    "        prompt = DEFAULT_JSONALYZE_PROMPT\n",
    "\n",
    "        # Instantiate in-memory SQLite database\n",
    "        db = sqlite_utils.Database(memory=True)\n",
    "        try:\n",
    "            # Load list of dictionaries into SQLite database\n",
    "            db[ev.table_name].insert_all(list_of_dict)\n",
    "        except sqlite_utils.utils.sqlite3.IntegrityError as exc:\n",
    "            print_text(\n",
    "                f\"Error inserting into table {table_name}, expected format:\"\n",
    "            )\n",
    "            print_text(\"[{col1: val1, col2: val2, ...}, ...]\")\n",
    "            raise ValueError(\"Invalid list_of_dict\") from exc\n",
    "\n",
    "        # Get the table schema\n",
    "        table_schema = db[table_name].columns_dict\n",
    "\n",
    "        # Get the SQL query with text-to-SQL prompt\n",
    "        response_str = await ev.llm.apredict(\n",
    "            prompt=prompt,\n",
    "            table_name=table_name,\n",
    "            table_schema=table_schema,\n",
    "            question=query,\n",
    "        )\n",
    "\n",
    "        sql_parser = DefaultSQLParser()\n",
    "\n",
    "        sql_query = sql_parser.parse_response_to_sql(response_str, ev.query)\n",
    "\n",
    "        try:\n",
    "            # Execute the SQL query\n",
    "            results = list(db.query(sql_query))\n",
    "        except sqlite_utils.utils.sqlite3.OperationalError as exc:\n",
    "            print_text(f\"Error executing query: {sql_query}\")\n",
    "            raise ValueError(\"Invalid query\") from exc\n",
    "\n",
    "        return JsonAnalyzerEvent(\n",
    "            sql_query=sql_query, table_schema=table_schema, results=results\n",
    "        )\n",
    "\n",
    "    @step\n",
    "    async def synthesize(\n",
    "        self, ctx: Context, ev: JsonAnalyzerEvent\n",
    "    ) -> StopEvent:\n",
    "        \"\"\"Synthesize the response.\"\"\"\n",
    "        llm = await ctx.store.get(\"llm\", default=None)\n",
    "        query = await ctx.store.get(\"query\", default=None)\n",
    "\n",
    "        response_str = llm.predict(\n",
    "            DEFAULT_RESPONSE_SYNTHESIS_PROMPT,\n",
    "            sql_query=ev.sql_query,\n",
    "            table_schema=ev.table_schema,\n",
    "            sql_response=ev.results,\n",
    "            query_str=query,\n",
    "        )\n",
    "\n",
    "        response_metadata = {\n",
    "            \"sql_query\": ev.sql_query,\n",
    "            \"table_schema\": str(ev.table_schema),\n",
    "        }\n",
    "\n",
    "        response = Response(response=response_str, metadata=response_metadata)\n",
    "\n",
    "        return StopEvent(result=response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Json List"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "json_list = [\n",
    "    {\n",
    "        \"name\": \"John Doe\",\n",
    "        \"age\": 25,\n",
    "        \"major\": \"Computer Science\",\n",
    "        \"email\": \"john.doe@example.com\",\n",
    "        \"address\": \"123 Main St\",\n",
    "        \"city\": \"New York\",\n",
    "        \"state\": \"NY\",\n",
    "        \"country\": \"USA\",\n",
    "        \"phone\": \"+1 123-456-7890\",\n",
    "        \"occupation\": \"Software Engineer\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Jane Smith\",\n",
    "        \"age\": 30,\n",
    "        \"major\": \"Business Administration\",\n",
    "        \"email\": \"jane.smith@example.com\",\n",
    "        \"address\": \"456 Elm St\",\n",
    "        \"city\": \"San Francisco\",\n",
    "        \"state\": \"CA\",\n",
    "        \"country\": \"USA\",\n",
    "        \"phone\": \"+1 234-567-8901\",\n",
    "        \"occupation\": \"Marketing Manager\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Michael Johnson\",\n",
    "        \"age\": 35,\n",
    "        \"major\": \"Finance\",\n",
    "        \"email\": \"michael.johnson@example.com\",\n",
    "        \"address\": \"789 Oak Ave\",\n",
    "        \"city\": \"Chicago\",\n",
    "        \"state\": \"IL\",\n",
    "        \"country\": \"USA\",\n",
    "        \"phone\": \"+1 345-678-9012\",\n",
    "        \"occupation\": \"Financial Analyst\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Emily Davis\",\n",
    "        \"age\": 28,\n",
    "        \"major\": \"Psychology\",\n",
    "        \"email\": \"emily.davis@example.com\",\n",
    "        \"address\": \"234 Pine St\",\n",
    "        \"city\": \"Los Angeles\",\n",
    "        \"state\": \"CA\",\n",
    "        \"country\": \"USA\",\n",
    "        \"phone\": \"+1 456-789-0123\",\n",
    "        \"occupation\": \"Psychologist\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Alex Johnson\",\n",
    "        \"age\": 27,\n",
    "        \"major\": \"Engineering\",\n",
    "        \"email\": \"alex.johnson@example.com\",\n",
    "        \"address\": \"567 Cedar Ln\",\n",
    "        \"city\": \"Seattle\",\n",
    "        \"state\": \"WA\",\n",
    "        \"country\": \"USA\",\n",
    "        \"phone\": \"+1 567-890-1234\",\n",
    "        \"occupation\": \"Civil Engineer\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Jessica Williams\",\n",
    "        \"age\": 32,\n",
    "        \"major\": \"Biology\",\n",
    "        \"email\": \"jessica.williams@example.com\",\n",
    "        \"address\": \"890 Walnut Ave\",\n",
    "        \"city\": \"Boston\",\n",
    "        \"state\": \"MA\",\n",
    "        \"country\": \"USA\",\n",
    "        \"phone\": \"+1 678-901-2345\",\n",
    "        \"occupation\": \"Biologist\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Matthew Brown\",\n",
    "        \"age\": 26,\n",
    "        \"major\": \"English Literature\",\n",
    "        \"email\": \"matthew.brown@example.com\",\n",
    "        \"address\": \"123 Peach St\",\n",
    "        \"city\": \"Atlanta\",\n",
    "        \"state\": \"GA\",\n",
    "        \"country\": \"USA\",\n",
    "        \"phone\": \"+1 789-012-3456\",\n",
    "        \"occupation\": \"Writer\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Olivia Wilson\",\n",
    "        \"age\": 29,\n",
    "        \"major\": \"Art\",\n",
    "        \"email\": \"olivia.wilson@example.com\",\n",
    "        \"address\": \"456 Plum Ave\",\n",
    "        \"city\": \"Miami\",\n",
    "        \"state\": \"FL\",\n",
    "        \"country\": \"USA\",\n",
    "        \"phone\": \"+1 890-123-4567\",\n",
    "        \"occupation\": \"Artist\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Daniel Thompson\",\n",
    "        \"age\": 31,\n",
    "        \"major\": \"Physics\",\n",
    "        \"email\": \"daniel.thompson@example.com\",\n",
    "        \"address\": \"789 Apple St\",\n",
    "        \"city\": \"Denver\",\n",
    "        \"state\": \"CO\",\n",
    "        \"country\": \"USA\",\n",
    "        \"phone\": \"+1 901-234-5678\",\n",
    "        \"occupation\": \"Physicist\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Sophia Clark\",\n",
    "        \"age\": 27,\n",
    "        \"major\": \"Sociology\",\n",
    "        \"email\": \"sophia.clark@example.com\",\n",
    "        \"address\": \"234 Orange Ln\",\n",
    "        \"city\": \"Austin\",\n",
    "        \"state\": \"TX\",\n",
    "        \"country\": \"USA\",\n",
    "        \"phone\": \"+1 012-345-6789\",\n",
    "        \"occupation\": \"Social Worker\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Christopher Lee\",\n",
    "        \"age\": 33,\n",
    "        \"major\": \"Chemistry\",\n",
    "        \"email\": \"christopher.lee@example.com\",\n",
    "        \"address\": \"567 Mango St\",\n",
    "        \"city\": \"San Diego\",\n",
    "        \"state\": \"CA\",\n",
    "        \"country\": \"USA\",\n",
    "        \"phone\": \"+1 123-456-7890\",\n",
    "        \"occupation\": \"Chemist\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Ava Green\",\n",
    "        \"age\": 28,\n",
    "        \"major\": \"History\",\n",
    "        \"email\": \"ava.green@example.com\",\n",
    "        \"address\": \"890 Cherry Ave\",\n",
    "        \"city\": \"Philadelphia\",\n",
    "        \"state\": \"PA\",\n",
    "        \"country\": \"USA\",\n",
    "        \"phone\": \"+1 234-567-8901\",\n",
    "        \"occupation\": \"Historian\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Ethan Anderson\",\n",
    "        \"age\": 30,\n",
    "        \"major\": \"Business\",\n",
    "        \"email\": \"ethan.anderson@example.com\",\n",
    "        \"address\": \"123 Lemon Ln\",\n",
    "        \"city\": \"Houston\",\n",
    "        \"state\": \"TX\",\n",
    "        \"country\": \"USA\",\n",
    "        \"phone\": \"+1 345-678-9012\",\n",
    "        \"occupation\": \"Entrepreneur\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Isabella Carter\",\n",
    "        \"age\": 28,\n",
    "        \"major\": \"Mathematics\",\n",
    "        \"email\": \"isabella.carter@example.com\",\n",
    "        \"address\": \"456 Grape St\",\n",
    "        \"city\": \"Phoenix\",\n",
    "        \"state\": \"AZ\",\n",
    "        \"country\": \"USA\",\n",
    "        \"phone\": \"+1 456-789-0123\",\n",
    "        \"occupation\": \"Mathematician\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Andrew Walker\",\n",
    "        \"age\": 32,\n",
    "        \"major\": \"Economics\",\n",
    "        \"email\": \"andrew.walker@example.com\",\n",
    "        \"address\": \"789 Berry Ave\",\n",
    "        \"city\": \"Portland\",\n",
    "        \"state\": \"OR\",\n",
    "        \"country\": \"USA\",\n",
    "        \"phone\": \"+1 567-890-1234\",\n",
    "        \"occupation\": \"Economist\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Mia Evans\",\n",
    "        \"age\": 29,\n",
    "        \"major\": \"Political Science\",\n",
    "        \"email\": \"mia.evans@example.com\",\n",
    "        \"address\": \"234 Lime St\",\n",
    "        \"city\": \"Washington\",\n",
    "        \"state\": \"DC\",\n",
    "        \"country\": \"USA\",\n",
    "        \"phone\": \"+1 678-901-2345\",\n",
    "        \"occupation\": \"Political Analyst\",\n",
    "    },\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(model=\"gpt-3.5-turbo\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run the Workflow!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w = JSONAnalyzeQueryEngineWorkflow()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "> Question: What is the maximum age among the individuals?"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "Answer: The maximum age among the individuals is 35."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Run a query\n",
    "\n",
    "query = \"What is the maximum age among the individuals?\"\n",
    "\n",
    "result = await w.run(\n",
    "    query=query, list_of_dict=json_list, llm=llm, table_name=DEFAULT_TABLE_NAME\n",
    ")\n",
    "\n",
    "display(\n",
    "    Markdown(\"> Question: {}\".format(query)),\n",
    "    Markdown(\"Answer: {}\".format(result)),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "> Question: How many individuals have an occupation related to science or engineering?"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "Answer: There are 0 individuals with an occupation related to science or engineering."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "query = \"How many individuals have an occupation related to science or engineering?\"\n",
    "\n",
    "result = await w.run(\n",
    "    query=query, list_of_dict=json_list, llm=llm, table_name=DEFAULT_TABLE_NAME\n",
    ")\n",
    "\n",
    "display(\n",
    "    Markdown(\"> Question: {}\".format(query)),\n",
    "    Markdown(\"Answer: {}\".format(result)),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "> Question: How many individuals have a phone number starting with '+1 234'?"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "Answer: There are 2 individuals with a phone number starting with '+1 234'."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "query = \"How many individuals have a phone number starting with '+1 234'?\"\n",
    "\n",
    "result = await w.run(\n",
    "    query=query, list_of_dict=json_list, llm=llm, table_name=DEFAULT_TABLE_NAME\n",
    ")\n",
    "\n",
    "display(\n",
    "    Markdown(\"> Question: {}\".format(query)),\n",
    "    Markdown(\"Answer: {}\".format(result)),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "> Question: What is the percentage of individuals residing in California (CA)?"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "Answer: The percentage of individuals residing in California (CA) is 18.75%."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "query = \"What is the percentage of individuals residing in California (CA)?\"\n",
    "\n",
    "result = await w.run(\n",
    "    query=query, list_of_dict=json_list, llm=llm, table_name=DEFAULT_TABLE_NAME\n",
    ")\n",
    "\n",
    "display(\n",
    "    Markdown(\"> Question: {}\".format(query)),\n",
    "    Markdown(\"Answer: {}\".format(result)),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "> Question: How many individuals have a major in Psychology?"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "Answer: There is 1 individual who has a major in Psychology."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "query = \"How many individuals have a major in Psychology?\"\n",
    "\n",
    "result = await w.run(\n",
    "    query=query, list_of_dict=json_list, llm=llm, table_name=DEFAULT_TABLE_NAME\n",
    ")\n",
    "\n",
    "display(\n",
    "    Markdown(\"> Question: {}\".format(query)),\n",
    "    Markdown(\"Answer: {}\".format(result)),\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llamaindex",
   "language": "python",
   "name": "llamaindex"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
